package com.mc.book.utils.sse;

import com.mc.book.pojo.Result;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import javax.servlet.http.HttpServletResponse;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;

/**
 *
 * @Author: dong
 * @Date: 2024/7/2 14:23
 */
@Slf4j
public class BookListener extends EventSourceListener {

    @Getter
    private CountDownLatch countDownLatch = new CountDownLatch(1);
    private HttpServletResponse rp;

    private StringBuffer output = new StringBuffer();

    public BookListener(HttpServletResponse rp) {
        this.rp = rp;
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        if (rp != null) {
            rp.setContentType("text/event-stream");
            rp.setCharacterEncoding("UTF-8");
            rp.setHeader("Cache-Control", "no-cache");
            rp.setStatus(200);
        }
    }

    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        log.info("sse连接关闭:{}");
        log.info("结果输出:{}" + output.toString());
        countDownLatch.countDown();
    }

    @Override
    public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
        try {
            output.append(data);
            if ("finish".equals(type)) {
//                log.info("请求结束{} {}", chatGlmDto.getMessageId(), output.toString());
            }
            if ("error".equals(type)) {
//                log.info("{}: {}source {}", chatGlmDto.getMessageId(), data, JSON.toJSONString(chatGlmDto));
            }
            if (rp != null) {
                if ("\n".equals(data)) {
//                    rp.getWriter().write("event:" + type + "\n");
//                    rp.getWriter().write("id:" + /*chatGlmDto.getMessageId() +*/ "\n");
                    rp.getWriter().write("data:\n\n");
                    rp.getWriter().flush();
                } else {
                    String[] dataArr = data.split("\\n");
                    for (int i = 0; i < dataArr.length; i++) {
//                        if (i == 0) {
//                            rp.getWriter().write("event:" + type + "\n");
//                            rp.getWriter().write("id:" + /*chatGlmDto.getMessageId() +*/ "\n");
//                        }
                        rp.getWriter().write("event: message\n");
                        if (i == dataArr.length - 1) {
                            rp.getWriter().write("data:"  + dataArr[i] + "\n\n");
//                            rp.getWriter().write("{\"code\":0,\"data\":" + dataArr[i] + "}\n\n");
                            rp.getWriter().flush();
                        } else {
                            rp.getWriter().write("data:"  + dataArr[i] + "\n\n");
//                            rp.getWriter().write("{\"code\":0,\"data\":"  + dataArr[i] + "}\n");
                            rp.getWriter().flush();
                        }
                    }
                }

            }
        } catch (Exception e) {
            log.error("消息错误[" + "]", e);
            countDownLatch.countDown();
            throw new RuntimeException(e);
        }
    }

    @Override
    public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
        countDownLatch.countDown();
        super.onFailure(eventSource, t, response);
    }
}
